Performance Improvements in Julia (Part 1): Lazy Broadcasting

Introduction

Readability and performance are two crucial features in coding. However, they're not always compatible, leading you to prioritize one over the other. In this post, we'll explore a tool that enables you to enhance performance without compromising readability: lazy broadcasting.

In programming, the term lazy refers to postponing the execution of some code until the result is required. The opposite is eager, which means executing the code as soon as the computer reads it. When it comes to broadcasting, Julia performs it eagerly by default. This post will show how lazy broadcasting can enhance performance.

Lazy broadcasting is useful for computing an output that requires several intermediate steps, as long as intermediate results are unimportant and we don't need to keep them. This improves performance by i) avoiding memory allocations for intermediate results, [note] Strictly speaking, we're referring to heap allocations. The concept involves several technical aspects that are not necessary for our purposes. See here if you want to know why this matter for performance, and here for an intuitive explanation of heap and stack allocations. and ii) combining operations into a single one, which enables Julia to perform further optimizations.

Comparing Approaches

The following packages are required for the code in this post. Note in particular the use of LazyArrays, which is what allows us to use lazy broadcasting.

using DataFrames, Distributions, Random, BenchmarkTools
using LazyArrays

The code below generates mock data for prices. It additionally creates functions in Julia representing equations (1), (2), (3), and (4).

Random.seed!(123)        # setting the seed for reproducibility
price_vector             = 2 .* rand(1_000)

quantity(price)          = 10 - 2.1 * price
revenue(quantity, price) = quantity * price
cost(quantity)           = 1.1 * quantity
profit(revenue, cost)    = revenue - cost

We now compare the performance of two broadcasting approaches. The first one is the default approach, which performs broadcasting eagerly. The second one performs the broadcast lazily, which we indicate by adding the macro @~. [note] The benchmarks in this post wrap each variable into the function ref(x) = (Ref(x))[]. This prevents the computer from realizing that we're repeating the same calculation. Otherwise, @btime would directly repeat stating the result after a few iterations rather than computing it, thus distorting the comparison.

function compute_profit1(price)
    compute_quantity = quantity.(price)
    compute_revenue  = revenue.(compute_quantity, price)
    compute_cost     = cost.(compute_quantity)
    
    profit.(compute_revenue, compute_cost)     
end

@btime compute_profit1(ref($price_vector));

  1.070 μs (4 allocations: 31.75 KiB)

function compute_profit2(price)
    compute_quantity = @~ quantity.(price)
    compute_revenue  = @~ revenue.(compute_quantity, price)
    compute_cost     = @~ cost.(compute_quantity)
    
    profit.(compute_revenue, compute_cost)
end

@btime compute_profit2(ref($price_vector));

  233.333 ns (1 allocation: 7.94 KiB)

As 1 μs is 1,000 ns, the lazy approach is faster. This occurs because the default broadcasting performs eager computations. Thus, compute_profit1 stores the results for quantities, revenue, and cost, and then computes the profit with those results. The procedure ends up performing 4 allocations: three for the intermediate steps, and one for the output of interest (i.e., profit).

In contrast, compute_profit2 indicates Julia that quantities, revenue, and cost should be calculated lazily. It's like saying to Julia: "do not compute anything at this point, let me only describe the operation involved". This allows Julia to avoid the allocations for the intermediate results, fuse all the operations into a single one, and provide the profit as the outcome. This explains why we've reduced 4 allocations to 1, where the only allocation is storing the vector of profits and hence is unavoidable.

Remark

Notice that the allocations would've occurred even if we hadn't stored the results in variables. This is because any expression that relies on broadcasting allocates, even if the final result is a scalar (which does not allocate). The following example demonstrates this.

Eager Broadcast (default) Lazy Broadcast

x        = [2,3,4]
eager(x) = sum(x .* x .* log.(x))

@btime eager(ref($x))

Output in REPL

  30.955 ns (1 allocation: 80 bytes)

x       = [2,3,4]
lazy(x) = sum(@~(x .* x .* log.(x)))

@btime lazy(ref($x))

Output in REPL

  29.819 ns (0 allocations: 0 bytes)